import numpy as np
from numpy.core.arrayprint import dtype_is_implied
from numpy.lib.function_base import extract
from numpy.lib.utils import source

class solution_FE_solver_1D_Possion:
    def __init__(self,a,b,number_of_elements):
        self.a = a
        self.b = b
        self.number_of_elements = number_of_elements

        

    def generate_PT(self):
        a = self.a
        b = self.b
        number_of_elements = self.number_of_elements

        P = np.linspace(a,b,number_of_elements+1)
        T = np.zeros((number_of_elements,2),dtype=int)
        T[:,0] = np.arange(number_of_elements)
        T[:,1] = T[:,0]+1

        return P,T


    def generate_PbTb_trial(self,basis_type_trial):
        if basis_type_trial==101:
            Pb_trial,Tb_trial = self.generate_PT()

        #elif basis_type_trial==102:

        return Pb_trial,Tb_trial

        

    def generate_PbTb_test(self,basis_type_test):
        if basis_type_test==101:
            Pb_test,Tb_test = self.generate_PT()

        #elif basis_type_test==102:

        return Pb_test,Tb_test

        



    def generate_boundarynodes(self,basis_type_trial):
        #boundarynodes[i,0] 表示第i个边界点的边界类型:
        #0-- dirichlet, 1--neumann, 2--robin
        #boundarynodes[i,1] 表示第i个边界点的全局坐标
        boundarynodes = np.zeros((2,2),dtype=int)

        boundarynodes[0,0] = 0
        boundarynodes[0,1] = 0

        boundarynodes[1,0] = 0

        number_of_elements = self.number_of_elements
        if basis_type_trial == 101:
            boundarynodes[1,1] = number_of_elements

        return boundarynodes



    def A_assemble_matrix_1D(self,P,T,Gauss_type,coe_fun,
                            Tb_trial, basis_type_trial,basis_der_x_trial,
                            Tb_test,  basis_type_test, basis_der_x_test):


        number_of_global_basis_fun_trial = np.max(Tb_trial) + 1
        number_of_global_basis_fun_test  = np.max(Tb_test) + 1

        number_of_local_basis_fun_trial = Tb_trial.shape[1]
        number_of_locla_basis_fun_test  = Tb_test.shape[1]

        number_of_elements = T.shape[0]




        #from scipy.sparse import csr_matrix
        #A = csr_matrix((number_of_global_basis_fun_test,number_of_global_basis_fun_trial))
        A = np.zeros((number_of_global_basis_fun_trial,number_of_global_basis_fun_test),dtype=float)

        for n in range(number_of_elements):

            vertices = P[T[n]]
            Gauss_weights,Gauss_nodes = self.generate_Gauss_local(vertices,Gauss_type)

            for alpha in range(number_of_local_basis_fun_trial):

                for beta in range(number_of_locla_basis_fun_test):
                    
                    basis_index_trial = alpha
                    basis_index_test = beta

                    int_value=self.Gauss_quad_1D_trial_test(coe_fun,Gauss_weights,Gauss_nodes,vertices,
                                                        basis_type_trial,basis_index_trial,basis_der_x_trial,
                                                        basis_type_test,basis_index_test,basis_der_x_test)

                    #S[beta,alpha]=int_value #单刚
                    #i = Tb_test[beta,n]
                    #j = Tb_trial[alpha,n]
                    A[Tb_test[n,beta],Tb_trial[n,beta]]=A[Tb_test[n,beta],Tb_trial[n,beta]]+int_value

        
        print(A)
        return A





    def b_assemble_vector_1D(self,P,T,Gauss_type,source_fun,
                                Tb_test,basis_type_test,basis_der_x_test_right):

        number_of_elements = T.shape[0] #number of cells
        number_of_locla_basis_fun_test = Tb_test.shape[1] #number of local test basis function
        number_of_global_basis_fun_test = np.max(Tb_test) + 1
        b = np.zeros(number_of_global_basis_fun_test,dtype=float)


        for n in range(number_of_elements):
            vertices = P[T[n]]
            Gauss_weights,Gauss_nodes = self.generate_Gauss_local(vertices,Gauss_type)

            for beta in range(number_of_locla_basis_fun_test):

                basis_index_test = beta

                int_value = self.Gauss_quad_1D_right_test(source_fun,Gauss_weights,Gauss_nodes,vertices,
                                                        basis_type_test,basis_index_test,basis_der_x_test_right)

                b[Tb_test[n,beta]]+=int_value


        return b





    def treat_Dirichler_boundary(self,basis_type_trial,Pb_trial,A,b):
        boundarynodes = self.generate_boundarynodes(basis_type_trial)
        nbn = boundarynodes.shape[0]

        for k in range(nbn):
            if boundarynodes[k,0] == 0:
                i = boundarynodes[k,1]
                A[i] = 0
                A[i,i] = 1
                b[i] = self.function_gb(Pb_trial[i])

        return A, b

    def solution(self,A,b):
        from scipy.sparse.linalg import spsolve
        return spsolve(A,b)



    def compute_max_FE_nodes(self,Pb_trial):
        error = 0
        return error



    def Gauss_quad_1D_trial_test(self,coe_fun,Gauss_weights,Gauss_nodes,vertices,
                                    basis_type_trial,basis_index_trial,basis_der_x_trial,
                                    basis_type_test, basis_index_test, basis_der_x_test):

        int_value=0
        Gpn = Gauss_weights.shape[0]

        for k in range(Gpn):
            int_value=int_value+Gauss_weights[k]*(coe_fun(Gauss_nodes[k])
                                        *self.FE_basis_local_fun_1D(Gauss_nodes[k],vertices,basis_type_trial,basis_index_trial,basis_der_x_trial)
                                        *self.FE_basis_local_fun_1D(Gauss_nodes[k],vertices,basis_type_test, basis_index_test, basis_der_x_test))

                

        return int_value

    def Gauss_quad_1D_right_test(self,source_fun,Gauss_weights,Gauss_nodes,vertices,
                                    basis_type_test,basis_index_test,basis_der_x_test_right):
        int_value = 0
        Gpn = Gauss_weights.shape[0]

        for k in range(Gpn):
            int_value=int_value+Gauss_weights[k]*(source_fun(Gauss_nodes[k])
                                        *self.FE_basis_local_fun_1D(Gauss_nodes[k],vertices,basis_type_test, basis_index_test, basis_der_x_test))
        
        return int_value

    def function_c(self,x):
        result = np.exp(x)
        return result

    def function_f(self,x):
        result = -np.exp(x)*(np.cos(x)-2.0*np.sin(x)-x*np.cos(x)-x*np.sin(x))
        return result

    def function_gb(self,x):
        a = self.a
        b = self.b
        if np.abs(x-a)<1e-15:
            result = 0.0
        elif np.abs(x-b)<1e-15:
            result = np.cos(b)
        else:
            ValueError('No boundary points!')

        return result

    def function_exactu(self,x):
        result = x*np.cos(x)
        return result



    def FE_basis_local_fun_1D(self,x,vertices,basis_type,basis_index,basis_der_x):
        #101: 1D linear
        #102: 1D quatratic

        h = vertices[1] - vertices[0]
        if basis_type==101:

            if basis_index==0:

                if basis_der_x==0:
                    result = (vertices[1]-x)/h

                elif basis_der_x==1:
                    result = -1.0/h

                elif basis_der_x>=2&type(basis_der_x)==int:
                    result = 0.0

                else:
                    ValueError('wrong, my curent code can not hanlde this derivative')



            elif basis_index==1:

                if basis_der_x==0:
                    result = (x-vertices[0])/h

                elif basis_der_x==1:
                    result = 1.0/h

                elif basis_der_x>=2&type(basis_der_x)==int:
                    result = 0.0

                else:
                    ValueError('wrong, my curent code can not hanlde this derivative')



            else:
                ValueError('wrong input for basis index! ')

        #elif basis_type==102:

        return result


    def generate_Gauss_local(self,vertices,Gauss_type):
        #此处是Gauss-Legrendre积分
        volume = vertices[1]-vertices[0]
        if Gauss_type == 1: #中心点积分，1阶段精度
            Gauss_weights = np.array([1.0],dtype=float)*volume

            Gauss_nodes = 0.5*(vertices[0]+vertices[1])

        elif Gauss_type == 2: #两点积分，3阶精度
            Gauss_weights = np.array([0.5,0.5],dtype=float)*volume

            lam = np.array([-0.57735026918963,0.57735026918963],dtype=float)
            Gauss_nodes = 0.5*(vertices[0]+vertices[1])+0.5*(vertices[1]-vertices[0])*lam



        elif Gauss_type == 3: #三点积分，5阶精度
            Gauss_weights = np.array([0.277777777777775,0.444444444444444,0.277777777777775],dtype=float)*volume

            lam = np.array([-0.77459666924148,0.0,0.77459666924148],dtype=float)
            Gauss_nodes = 0.5*(vertices[0]+vertices[1])+0.5*(vertices[1]-vertices[0])*lam

        elif Gauss_type == 4: #4点积分，7阶精度
            Gauss_weights = np.array([0.173927422568725,0.326072577431275,0.326072577431275,0.173927422568725])*volume

            lam = np.array([-0.86113631159405,-0.33998104358486,0.33998104358486,0.86113631159405],dtype=float)
            Gauss_nodes = 0.5*(vertices[0]+vertices[1])+0.5*(vertices[1]-vertices[0])*lam

        return Gauss_weights,Gauss_nodes





if __name__ == '__main__':
    a = 0.0
    b = 1.0
    number_of_elements = 2
    basis_type_trial = 101
    basis_type_test = 101
    basis_der_x_trial = 1
    basis_der_x_test = 1
    basis_der_x_test_right = 0
    Gauss_type = 2

    onedpossion = solution_FE_solver_1D_Possion(a,b,number_of_elements)

    coe_fun = onedpossion.function_c
    source_fun = onedpossion.function_f


    P,T = onedpossion.generate_PT()
    Pb_trial,Tb_trail = onedpossion.generate_PbTb_trial(basis_type_trial)
    Pb_test,Tb_test   = onedpossion.generate_PbTb_test(basis_type_test)
    boundarynodes = onedpossion.generate_boundarynodes(basis_type_trial)


    A = onedpossion.A_assemble_matrix_1D(P,T,Gauss_type,coe_fun,
                                        Tb_trail,basis_type_trial,basis_der_x_trial,
                                        Tb_test, basis_type_test, basis_der_x_test)
    #print(A)

    b = onedpossion.b_assemble_vector_1D(P,T,Gauss_type,source_fun,
                                        Tb_test,basis_type_test,basis_der_x_test_right)


    A, b = onedpossion.treat_Dirichler_boundary(basis_type_trial,Pb_trial,A,b)

    uh = onedpossion.solution(A,b)

